import pandas as pd
from sklearn.model_selection import GroupKFold
from xgboost import XGBRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error

# 读取数据
data_path = 'D:\\学习&科研\\华为手表项目\\华为数据\\试验记录表\\all_stages_df_statistics.csv'
df = pd.read_csv(data_path)

# 将 polar_hr 和 polar_rr 列转换为适合模型的格式，例如取平均值
df['polar_hr'] = df['polar_hr'].apply(lambda x: eval(x)[0] if isinstance(eval(x), list) and len(eval(x)) > 0 else 0)
df['polar_rr'] = df['polar_rr'].apply(lambda x: eval(x)[0] if isinstance(eval(x), list) and len(eval(x)) > 0 else 0)

# 筛选状态为 running 的数据
df = df[df['state'] == 'running']

# 初始化 GroupKFold
gkf = GroupKFold(n_splits=5)

# 选择特征和目标变量
X = df[['polar_hr_mean', 'speed',  'polar_hr_min', 'polar_hr_max',
         'polar_rr_mean',  'polar_rr_min']]
y = df['physiology_RPE']
groups = df['number']  # 用于分组的列

# 初始化结果列表
results = []
all_y_true = []
all_y_pred = []

# 使用 GroupKFold 进行分组划分
for train_index, test_index in gkf.split(X, y, groups=groups):
    X_train, X_test = X.iloc[train_index], X.iloc[test_index]
    y_train, y_test = y.iloc[train_index], y.iloc[test_index]

    # 创建 XGBoost 回归模型
    model = XGBRegressor(n_estimators=100, random_state=42)

    # 训练模型
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # 评估模型
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)

    # 保存每个分组的结果
    results.append({
        'Mean Squared Error': mse,
        'R² Score': r2,
        'Mean Absolute Error': mae
    })

    # 保存所有的真实值和预测值
    all_y_true.extend(y_test)
    all_y_pred.extend(y_pred)

# 将结果转换为 DataFrame 并输出
results_df = pd.DataFrame(results)
print(results_df)

# 计算整体模型性能
overall_mse = mean_squared_error(all_y_true, all_y_pred)
overall_r2 = r2_score(all_y_true, all_y_pred)
overall_mae = mean_absolute_error(all_y_true, all_y_pred)

# 输出整体模型性能
print(f'Overall Mean Squared Error: {overall_mse}')
print(f'Overall R² Score: {overall_r2}')
print(f'Overall Mean Absolute Error: {overall_mae}')
